import torch
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List
from train_eval.initialization import initialize_prediction_model, initialize_dataset, get_specific_args
from nuscenes.eval.prediction.splits import get_prediction_challenge_split
from nuscenes.prediction.input_representation.static_layers import StaticLayerRasterizer
from nuscenes.prediction.input_representation.agents import AgentBoxesWithFadedHistory
from nuscenes.prediction.input_representation.interface import InputRepresentation
from nuscenes.prediction.input_representation.combinators import Rasterizer
import train_eval.utils as u
import imageio
import os


# Initialize device:
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class Visualizer:
    """
    Class for visualizing predictions generated by trained model
    """
    def __init__(self, cfg: Dict, data_root: str, data_dir: str, checkpoint_path: str):
        """
        Initialize evaluator object
        :param cfg: Configuration parameters
        :param data_root: Root directory with data
        :param data_dir: Directory with extracted, pre-processed data
        :param checkpoint_path: Path to checkpoint with trained weights
        """

        # Initialize dataset
        ds_type = cfg['dataset'] + '_' + cfg['agent_setting'] + '_' + cfg['input_representation']
        spec_args = get_specific_args(cfg['dataset'], data_root, cfg['version'] if 'version' in cfg.keys() else None)
        test_set = initialize_dataset(ds_type, ['load_data', data_dir, cfg['test_set_args']] + spec_args)
        self.ds = test_set

        # Initialize model
        self.model = initialize_prediction_model(cfg['encoder_type'], cfg['aggregator_type'], cfg['decoder_type'],
                                                 cfg['encoder_args'], cfg['aggregator_args'], cfg['decoder_args'])
        self.model = self.model.float().to(device)
        self.model.eval()

        # Load checkpoint
        checkpoint = torch.load(checkpoint_path)
        self.model.load_state_dict(checkpoint['model_state_dict'])

    def visualize(self, output_dir: str, dataset_type: str):
        """
        Generate visualizations for predictions
        :param output_dir: results directory to dump visualizations in
        :param dataset_type: e.g. 'nuScenes'. Visualizations will vary based on dataset.
        :return:
        """
        if dataset_type == 'nuScenes':
            self.visualize_nuscenes(output_dir)

    def visualize_nuscenes(self, output_dir):
        index_list = self.get_vis_idcs_nuscenes()
        if not os.path.isdir(os.path.join(output_dir, 'results', 'gifs')):
            os.mkdir(os.path.join(output_dir, 'results', 'gifs'))
        for n, indices in enumerate(index_list):
            imgs = self.generate_nuscenes_gif(indices)
            filename = os.path.join(output_dir, 'results', 'gifs', 'example' + str(n) + '.gif')
            imageio.mimsave(filename, imgs, format='GIF', fps=2)

    def get_vis_idcs_nuscenes(self):
        """
        Returns list of list of indices for generating gifs for the nuScenes val set.
        Instance tokens are hardcoded right now. I'll fix this later (TODO)
        """
        token_list = get_prediction_challenge_split('val', dataroot=self.ds.helper.data.dataroot)
        instance_tokens = [token_list[idx].split("_")[0] for idx in range(len(token_list))]
        unique_instance_tokens = []
        for i_t in instance_tokens:
            if i_t not in unique_instance_tokens:
                unique_instance_tokens.append(i_t)

        instance_tokens_to_visualize = [54, 98, 91, 5, 114, 144, 291, 204, 312, 187, 36, 267, 146]

        idcs = []
        for i_t_id in instance_tokens_to_visualize:
            idcs_i_t = [i for i in range(len(instance_tokens)) if instance_tokens[i] == unique_instance_tokens[i_t_id]]
            idcs.append(idcs_i_t)

        return idcs

    def generate_nuscenes_gif(self, idcs: List[int]):
        """
        Generates gif of predictions for the given set of indices.
        :param idcs: val set indices corresponding to a particular instance token.
        """

        # Raster maps for visualization.
        map_extent = self.ds.map_extent
        resolution = 0.1
        static_layer_rasterizer = StaticLayerRasterizer(self.ds.helper,
                                                        resolution=resolution,
                                                        meters_ahead=map_extent[3],
                                                        meters_behind=-map_extent[2],
                                                        meters_left=-map_extent[0],
                                                        meters_right=map_extent[1])

        agent_rasterizer = AgentBoxesWithFadedHistory(self.ds.helper, seconds_of_history=1,
                                                      resolution=resolution,
                                                      meters_ahead=map_extent[3],
                                                      meters_behind=-map_extent[2],
                                                      meters_left=-map_extent[0],
                                                      meters_right=map_extent[1])

        raster_maps = InputRepresentation(static_layer_rasterizer, agent_rasterizer, Rasterizer())

        imgs = []
        for idx in idcs:

            # Load data
            data = self.ds[idx]
            data = u.send_to_device(u.convert_double_to_float(u.convert2tensors(data)))
            i_t = data['inputs']['instance_token']
            s_t = data['inputs']['sample_token']

            # Get raster map
            hd_map = raster_maps.make_input_representation(i_t, s_t)
            r, g, b = hd_map[:, :, 0] / 255, hd_map[:, :, 1] / 255, hd_map[:, :, 2] / 255
            hd_map_gray = 0.2989 * r + 0.5870 * g + 0.1140 * b

            # Predict
            predictions = self.model(data['inputs'])

            # Plot
            fig, ax = plt.subplots(1, 3, figsize=(15, 5))
            ax[0].imshow(hd_map, extent=self.ds.map_extent)
            ax[1].imshow(hd_map_gray, cmap='gist_gray', extent=self.ds.map_extent)
            ax[2].imshow(hd_map_gray, cmap='gist_gray', extent=self.ds.map_extent)

            for n, traj in enumerate(predictions['traj'][0]):
                ax[1].plot(traj[:, 0].detach().cpu().numpy(), traj[:, 1].detach().cpu().numpy(), lw=4,
                           color='r', alpha=0.8)
                ax[1].scatter(traj[-1, 0].detach().cpu().numpy(), traj[-1, 1].detach().cpu().numpy(), 60,
                              color='r', alpha=0.8)

            traj_gt = data['ground_truth']['traj'][0]
            ax[2].plot(traj_gt[:, 0].detach().cpu().numpy(), traj_gt[:, 1].detach().cpu().numpy(), lw=4, color='g')
            ax[2].scatter(traj_gt[-1, 0].detach().cpu().numpy(), traj_gt[-1, 1].detach().cpu().numpy(), 60, color='g')

            ax[0].axis('off')
            ax[1].axis('off')
            ax[2].axis('off')
            fig.tight_layout(pad=0)
            ax[0].margins(0)
            ax[1].margins(0)
            ax[2].margins(0)

            fig.canvas.draw()
            image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
            image_from_plot = image_from_plot.reshape(fig.canvas.get_width_height()[::-1] + (3,))
            imgs.append(image_from_plot)
            plt.close(fig)

        return imgs
